package com.gitee.wsl.security.hash

import com.gitee.wsl.ext.base.Arrays.arraycopy
import com.gitee.wsl.security.digest.Digest
import com.gitee.wsl.security.digest.Mac


class HMAC {

    companion object {
        fun hmacSHA1(key: ByteArray, data: ByteArray): Hash = hmac(key, data, SHA1())

        fun hmacSHA256(key: ByteArray, data: ByteArray): Hash = hmac(key, data, SHA256())

        fun hmacSHA512(key: ByteArray, data: ByteArray): Hash = hmac(key, data, SHA512())

        fun hmacMD5(key: ByteArray, data: ByteArray): Hash = hmac(key, data, MD5())

        internal fun hmac(key: ByteArray, data: ByteArray, hasher: Hasher): Hash {
            var key = key
            val blockSize = hasher.chunkSize
            if (key.size > blockSize) {
                hasher.reset()
                hasher.update(key)
                key = hasher.digest().bytes
            }
            if (key.size < blockSize) {
                val newKey = ByteArray(blockSize)
                arraycopy(key, 0, newKey, 0, key.size)
                key = newKey
            }

            val oKeyPad = ByteArray(blockSize) { (0x5c xor key[it].toInt()).toByte() }
            val iKeyPad = ByteArray(blockSize) { (0x36 xor key[it].toInt()).toByte() }

            hasher.reset()
            hasher.update(iKeyPad)
            hasher.update(data)
            val h1 = hasher.digest().bytes

            hasher.reset()
            hasher.update(oKeyPad)
            hasher.update(h1)
            return hasher.digest()
        }
    }
}


fun ByteArray.hmacSha1(key: ByteArray) = HMAC.hmacSHA1(key,this)

fun ByteArray.hmacSha256(key: ByteArray) = HMAC.hmacSHA256(key,this)

fun ByteArray.hmacSha512(key: ByteArray) = HMAC.hmacSHA512(key,this)

fun ByteArray.hmacMD5(key: ByteArray) = HMAC.hmacMD5(key,this)

class HMacImp(
    private val digest: Digest,
    key: ByteArray
) : Mac {
    override val macSize: Int
        get() = digest.digestSize
    override val blockSize: Int
        get() = digest.blockSize
    override val algorithmName: String
        get() = "${digest.algorithmName}/HMAC"

    private val inputPad = ByteArray(blockSize)
    private val outputBuf = ByteArray(blockSize + macSize)

    init {
        digest.reset()
        var keyLength = key.size
        if (keyLength > blockSize) {
            digest.update(key)
            digest.digest(inputPad)
            keyLength = macSize
        } else {
            key.copyInto(inputPad)
        }

        inputPad.fill(0, keyLength)
        inputPad.copyInto(outputBuf)

        inputPad.indices.forEach {
            inputPad[it] = (inputPad[it].toInt() xor IPAD).toByte()
        }
        outputBuf.indices.forEach {
            outputBuf[it] = (outputBuf[it].toInt() xor OPAD).toByte()
        }
        digest.update(inputPad)
    }

    override fun reset() {
        digest.reset()
        digest.update(inputPad)
    }

    override fun digest(destination: ByteArray, destinationOffset: Int) {
        val blockSize = blockSize
        digest.digest(outputBuf, blockSize)
        digest.update(outputBuf)
        digest.digest(destination, destinationOffset)
        outputBuf.fill(0, blockSize)
        digest.update(inputPad)
    }

    override fun update(byte: Byte): HMacImp = apply {
        digest.update(byte)
    }

    override fun update(source: ByteArray, startIndex: Int, endIndex: Int): HMacImp = apply {
        digest.update(source, startIndex, endIndex)
    }

    public companion object {
        private const val IPAD = 0x36
        private const val OPAD = 0x5C
    }
}
